#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import time
import uuid
import shutil
from typing import Dict, Any, Optional, List
import tempfile
from pathlib import Path

from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from loguru import logger
from pydantic import BaseModel, ConfigDict, Field
from obs import ObsClient

# MinerU 相关导入
from mineru.data.dataset import PymuDocDataset
from mineru.model.doc_analyze import doc_analyze
from mineru.config.enums import SupportedPdfParseMethod
from mineru.utils.image_handler import ImageHandler

# 创建 FastAPI 应用
app = FastAPI(
    title="MinerU 服务",
    description="提供文档处理和分析服务",
    version="1.0.0",
)

# 配置 OBS 客户端
obs_client = ObsClient(
    access_key_id=os.getenv("OBS_ACCESS_KEY"),
    secret_access_key=os.getenv("OBS_SECRET_KEY"),
    server=os.getenv("OBS_ENDPOINT")
)
OBS_BUCKET = os.getenv("OBS_BUCKET")

# 数据模型
class ProcessRequest(BaseModel):
    """文档处理请求模型"""
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    task_id: str = Field(..., description="任务ID")
    image_handling: str = Field(default="obs", description="图片处理方式: 'obs'|'remove'|'keep'")
    options: Optional[Dict[str, Any]] = Field(default={}, description="处理选项")

class ProcessResponse(BaseModel):
    """文档处理响应模型"""
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    task_id: str = Field(..., description="任务ID")
    file_name: str = Field(..., description="原始文件名")
    markdown: str = Field(..., description="转换后的Markdown内容")
    process_time: float = Field(..., description="处理时间（秒）")
    engine_name: str = Field(..., description="使用的引擎名称")
    markdown_file: str = Field(..., description="Markdown文件路径")
    images: List[Dict[str, Any]] = Field(default=[], description="提取的图片信息")

# 状态记录
processing_tasks = {}

async def handle_images(
    images_dir: str,
    image_handling: str,
    task_id: str
) -> List[Dict[str, Any]]:
    """处理提取的图片
    
    Args:
        images_dir: 图片目录
        image_handling: 处理方式
        task_id: 任务ID
    
    Returns:
        处理后的图片信息列表
    """
    images = []
    if not os.path.exists(images_dir):
        return images
        
    for img_file in os.listdir(images_dir):
        if not img_file.lower().endswith(('.png', '.jpg', '.jpeg')):
            continue
            
        img_path = os.path.join(images_dir, img_file)
        
        if image_handling == "obs":
            # 上传到 OBS
            object_key = f"images/{task_id}/{img_file}"
            try:
                resp = obs_client.putFile(OBS_BUCKET, object_key, img_path)
                if resp.status < 300:
                    url = f"https://{OBS_BUCKET}.{os.getenv('OBS_ENDPOINT')}/{object_key}"
                    images.append({
                        "filename": img_file,
                        "format": Path(img_file).suffix.lstrip("."),
                        "url": url,
                        "status": "uploaded"
                    })
                else:
                    logger.error(f"上传图片失败: {resp.status} - {img_file}")
                    images.append({
                        "filename": img_file,
                        "format": Path(img_file).suffix.lstrip("."),
                        "error": f"上传失败: {resp.status}",
                        "status": "failed"
                    })
            except Exception as e:
                logger.exception(f"上传图片时出错: {str(e)}")
                images.append({
                    "filename": img_file,
                    "format": Path(img_file).suffix.lstrip("."),
                    "error": str(e),
                    "status": "failed"
                })
        elif image_handling == "keep":
            # 保留本地图片
            images.append({
                "filename": img_file,
                "format": Path(img_file).suffix.lstrip("."),
                "path": img_path,
                "status": "local"
            })
    
    return images

async def cleanup_temp_files(temp_dir: str):
    """清理临时文件
    
    Args:
        temp_dir: 临时目录路径
    """
    try:
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
    except Exception as e:
        logger.error(f"清理临时文件失败: {str(e)}")

@app.get("/")
async def root():
    return {"message": "MinerU 服务已启动"}

@app.get("/health")
async def health_check():
    """健康检查接口"""
    try:
        # 尝试导入必要的库
        from mineru.data.dataset import PymuDocDataset
        return {"status": "ok", "message": "服务运行正常"}
    except ImportError as e:
        return {"status": "error", "message": f"缺少必要的依赖: {str(e)}"}

@app.post("/process", response_model=ProcessResponse)
async def process_document(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    request: ProcessRequest = None,
):
    """处理文档
    
    使用 MinerU 引擎处理文档，支持以下功能：
    1. 文档转换为 Markdown
    2. 图片提取和处理（支持上传到 OBS 或移除）
    3. OCR 支持
    """
    if not request:
        raise HTTPException(status_code=400, detail="请求数据不完整")
    
    # 记录任务开始时间
    start_time = time.time()
    
    try:
        # 创建临时工作目录
        temp_dir = tempfile.mkdtemp()
        local_image_dir = os.path.join(temp_dir, "images")
        local_md_dir = os.path.join(temp_dir, "markdown")
        os.makedirs(local_image_dir, exist_ok=True)
        os.makedirs(local_md_dir, exist_ok=True)
        
        # 准备文件名
        name_without_suff = Path(file.filename).stem
        image_dir = "images"  # 相对路径，用于 Markdown 中的图片引用
        
        # 读取上传的文件内容
        content = await file.read()
        
        # 创建数据集实例
        ds = PymuDocDataset(content)
        
        # 根据文档类型选择处理方式
        if ds.classify() == SupportedPdfParseMethod.OCR:
            logger.info(f"使用 OCR 模式处理文档: {file.filename}")
            result = ds.process(
                doc_analyze,
                ocr=True,
                image_dir=local_image_dir,
                image_prefix=image_dir
            )
        else:
            logger.info(f"使用文本模式处理文档: {file.filename}")
            result = ds.process(
                doc_analyze,
                ocr=False,
                image_dir=local_image_dir,
                image_prefix=image_dir
            )
        
        # 保存 Markdown 文件
        markdown_file = os.path.join(local_md_dir, f"{name_without_suff}.md")
        with open(markdown_file, "w", encoding="utf-8") as f:
            f.write(result.markdown)
        
        # 处理图片
        images = await handle_images(
            local_image_dir,
            request.image_handling,
            request.task_id
        )
        
        # 计算处理时间
        process_time = time.time() - start_time
        
        # 创建响应
        response = ProcessResponse(
            task_id=request.task_id,
            file_name=file.filename,
            markdown=result.markdown,
            process_time=process_time,
            engine_name="MinerU",
            markdown_file=markdown_file,
            images=images
        )
        
        # 添加清理任务
        background_tasks.add_task(cleanup_temp_files, temp_dir)
        
        return response
        
    except Exception as e:
        logger.exception(f"处理文档时出错: {str(e)}")
        # 确保清理临时文件
        if 'temp_dir' in locals():
            background_tasks.add_task(cleanup_temp_files, temp_dir)
        raise HTTPException(status_code=500, detail=f"处理文档时出错: {str(e)}")

@app.get("/tasks/{task_id}")
async def get_task_status(task_id: str):
    """获取任务状态"""
    if task_id in processing_tasks:
        return processing_tasks[task_id]
    raise HTTPException(status_code=404, detail="任务不存在")

if __name__ == "__main__":
    import uvicorn
    
    # 从环境变量获取配置
    host = os.getenv("MINERU_SERVICE_HOST", "0.0.0.0")
    port = int(os.getenv("MINERU_SERVICE_PORT", "8001"))
    
    uvicorn.run(
        "main:app",
        host=host,
        port=port,
        reload=False  # 生产环境禁用热重载
    ) 